import os
import sys
import numpy as np
import h5py
import matplotlib.pyplot as plt
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))), 'util', 'file'))

class LungNoduleMalignancyDataset:
    """
    images and their labels, see https://gitee.com/yawei-gitee/medical-datasets-ongoing/tree/master/lung/lungNoduleMalignancy. 
    
    * Samples: 6691
    * Original label: 0/1
    * Task: binary classification
    """
    def __init__(self, givenDirOfDataset: str):
        self.__name = 'LNM'
        self.__dir = givenDirOfDataset
        self.__numberOfSamples = 6691
        self.__samples = None
        self.__labels = None 
        self.__randomSeed = 0
        self.__init()
        pass
    
    def __init(self):
        self.__initData()
        self.__shuffle()
        pass
    
    def __shuffle(self):
        self.__shuffleSamples()
        self.__shuffleLabels()
        pass
    
    def __shuffleSamples(self):
        """
        re-organize the samples randomly
        """
        np.random.seed(self.__randomSeed)
        self.__samples = np.random.permutation(self.__samples)
        pass
    
    def __shuffleLabels(self):
        """
        re-organize the samples randomly
        """
        np.random.seed(self.__randomSeed)
        self.__labels = np.random.permutation(self.__labels)
        pass

    def __initSamples(self):
        fullPath = f"{self.__dir}{os.sep}data.hdf5"
        with h5py.File(fullPath, 'r') as file:
            dataset = file['ct_slices']
            self.__samples = dataset[:]
        pass

    def __initLabels(self):
        fullPath = f"{self.__dir}{os.sep}data.hdf5"
        with h5py.File(fullPath, 'r') as file:
            score = file['slice_class']
            self.__labels = score[:]
        pass

    def __initData(self):
        self.__initSamples()
        self.__initLabels()
        pass

    def querySample(self, givenSampleId):
        return self.__samples[givenSampleId]
    
    def queryLabel(self, givenSampleId):
        return self.__labels[givenSampleId]
    
    def querySampleList(self, givenSampleIds: list):
        assert(isinstance(givenSampleIds, list))
        sampleList = [self.__samples[id] for id in givenSampleIds] 
        return np.array(sampleList)
    
    def getSamples(self):
        return self.__samples

    def getLabels(self):
        return self.__labels
    
    def showSample(self, givenId: int):
        plt.imshow(self.__samples[givenId], cmap='gray')
        plt.title(f'{givenId}-th sample')
        plt.show()
        pass

    @property
    def numberOfSamples(self):
        return self.__numberOfSamples

    @property
    def name(self):
        return self.__name
    pass










class Test:
    def execute(self):
        vsd = LungNoduleMalignancyDataset(givenDirOfDataset='E:\\images\\lung\\LungNoduleMalignancy\\basic')
        vsd.showSample(givenId=50)
        pass
    pass


if __name__ == '__main__':
    Test().execute()





